#include "remote_upgrade.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/stat.h>
#include <pthread.h>
#include <signal.h>
#include <curl/curl.h>

// 全局升级上下文
static upgrade_context_t g_upgrade_ctx = {0};
static pthread_t g_status_thread;
static volatile bool g_running = false;

// 函数声明
static void* status_report_thread(void* arg);
static int init_status_reporting(const char* status_topic);
static void stop_status_reporting(void);
static void update_status(upgrade_status_t status, const char* message);
static int init_mqtt_connection(const char* broker, int port,
                                const char* username, const char* password);
static char* calculate_file_md5(const char* filename);
static size_t write_data(void *ptr, size_t size, size_t nmemb, FILE *stream);
static int download_file_with_curl(const char* url, const char* output_path);
static int download_file_with_wget(const char* url, const char* output_path);
static int download_file(const char* url, const char* output_path);
static int execute_upgrade_script(const char* script_path);
static const char* get_upgrade_filename(const char* download_url);

// MQTT连接回调
static void mqtt_connect_cb(void *data)
{
    printf("MQTT connected successfully\n");
    (void)data;
    if (g_upgrade_ctx.mqtt_client) {
        g_upgrade_ctx.current_status = UPGRADE_STATUS_CONNECTING;
        g_upgrade_ctx.status_message = "MQTT connected, starting upgrade process";
    }
}

// MQTT断开回调
static void mqtt_disconnect_cb(void *data)
{
    printf("MQTT disconnected\n");
    (void)data;
}

// 状态上报线程
static void* status_report_thread(void* arg)
{
    upgrade_context_t* ctx = (upgrade_context_t*)arg;

    while (g_running && !ctx->stop_reporting) {
        // 构建状态消息
        char payload[512];
        const char* status_str = "";

        switch (ctx->current_status) {
        case UPGRADE_STATUS_IDLE: status_str = "idle"; break;
        case UPGRADE_STATUS_CONNECTING: status_str = "connecting"; break;
        case UPGRADE_STATUS_DOWNLOADING: status_str = "downloading"; break;
        case UPGRADE_STATUS_DOWNLOADED: status_str = "downloaded"; break;
        case UPGRADE_STATUS_VERIFYING: status_str = "verifying"; break;
        case UPGRADE_STATUS_VERIFIED: status_str = "verified"; break;
        case UPGRADE_STATUS_STARTING: status_str = "starting"; break;
        case UPGRADE_STATUS_RUNNING: status_str = "running"; break;
        case UPGRADE_STATUS_SUCCESS: status_str = "success"; break;
        case UPGRADE_STATUS_FAILED: status_str = "failed"; break;
        }

        snprintf(payload, sizeof(payload),
                 "{\"status\":\"%s\", \"message\":\"%s\", \"timestamp\":%ld}",
                 status_str, ctx->status_message, time(NULL));

        // 发布状态消息
        if (ctx->mqtt_client && ctx->status_topic) {
            neu_mqtt_client_publish(ctx->mqtt_client, 0, (char*)ctx->status_topic,
                                    (uint8_t*)payload, strlen(payload), NULL, NULL);
        }

        printf("Status: %s - %s\n", status_str, ctx->status_message);

        // 每秒上报一次
        sleep(1);
    }

    return NULL;
}

// 初始化状态上报
static int init_status_reporting(const char* status_topic)
{
    g_upgrade_ctx.status_topic = status_topic;
    g_upgrade_ctx.current_status = UPGRADE_STATUS_IDLE;
    g_upgrade_ctx.status_message = "Initializing";
    g_upgrade_ctx.stop_reporting = false;

    g_running = true;
    if (pthread_create(&g_status_thread, NULL, status_report_thread, &g_upgrade_ctx) != 0) {
        printf("Failed to create status report thread\n");
        return -1;
    }

    return 0;
}

// 停止状态上报
static void stop_status_reporting(void)
{
    g_upgrade_ctx.stop_reporting = true;
    g_running = false;

    if (g_status_thread) {
        pthread_join(g_status_thread, NULL);
    }
}

// 更新状态
static void update_status(upgrade_status_t status, const char* message)
{
    g_upgrade_ctx.current_status = status;
    g_upgrade_ctx.status_message = message;
}

// 初始化MQTT连接
static int init_mqtt_connection(const char* broker, int port,
                                const char* username, const char* password)
{
    // 创建MQTT客户端
    g_upgrade_ctx.mqtt_client = neu_mqtt_client_from_addr(broker, port, NEU_MQTT_VERSION_V311);
    if (!g_upgrade_ctx.mqtt_client) {
        printf("Failed to create MQTT client\n");
        return -1;
    }

    // 设置回调
    neu_mqtt_client_set_connect_cb(g_upgrade_ctx.mqtt_client, mqtt_connect_cb, NULL);
    neu_mqtt_client_set_disconnect_cb(g_upgrade_ctx.mqtt_client, mqtt_disconnect_cb, NULL);

    // 设置认证信息
    if (username && password) {
        neu_mqtt_client_set_user(g_upgrade_ctx.mqtt_client, username, password);
    }

    // 打开连接
    if (neu_mqtt_client_open(g_upgrade_ctx.mqtt_client) != 0) {
        printf("Failed to open MQTT connection\n");
        neu_mqtt_client_free(g_upgrade_ctx.mqtt_client);
        g_upgrade_ctx.mqtt_client = NULL;
        return -1;
    }

    // 等待连接建立
    int attempts = 0;
    while (!neu_mqtt_client_is_connected(g_upgrade_ctx.mqtt_client) && attempts < 50) {
        struct timespec t1 = { .tv_sec  = 100 / 1000,
                              .tv_nsec = 1000 * 1000 *
                                         (100 % 1000) };
        struct timespec t2 = { 0 };
        nanosleep(&t1,&t2);
        attempts++;
    }

    if (!neu_mqtt_client_is_connected(g_upgrade_ctx.mqtt_client)) {
        printf("MQTT connection timeout\n");
        neu_mqtt_client_close(g_upgrade_ctx.mqtt_client);
        neu_mqtt_client_free(g_upgrade_ctx.mqtt_client);
        g_upgrade_ctx.mqtt_client = NULL;
        return -1;
    }

    return 0;
}

// 计算文件MD5
static char* calculate_file_md5(const char* filename)
{
    static char md5[33] = {0};
    char command[512];

    snprintf(command, sizeof(command), "md5sum %s | cut -d' ' -f1", filename);
    FILE* fp = popen(command, "r");
    if (fp) {
        if (fgets(md5, sizeof(md5), fp)) {
            md5[strcspn(md5, "\n")] = 0;
        }
        pclose(fp);
    }

    return md5;
}

// 下载进度回调（用于CURL）
static size_t write_data(void *ptr, size_t size, size_t nmemb, FILE *stream)
{
    return fwrite(ptr, size, nmemb, stream);
}

// 使用CURL下载文件
static int download_file_with_curl(const char* url, const char* output_path)
{
    CURL *curl;
    FILE *fp;
    CURLcode res;

    curl = curl_easy_init();
    if (!curl) {
        printf("Failed to initialize CURL\n");
        return -1;
    }

    fp = fopen(output_path, "wb");
    if (!fp) {
        printf("Failed to create output file: %s\n", output_path);
        curl_easy_cleanup(curl);
        return -1;
    }

    curl_easy_setopt(curl, CURLOPT_URL, url);
    curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, write_data);
    curl_easy_setopt(curl, CURLOPT_WRITEDATA, fp);
    curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
    curl_easy_setopt(curl, CURLOPT_TIMEOUT, 300L); // 5分钟超时

    printf("Downloading: %s\n", url);
    update_status(UPGRADE_STATUS_DOWNLOADING, "Downloading upgrade package");

    res = curl_easy_perform(curl);

    fclose(fp);
    curl_easy_cleanup(curl);

    if (res != CURLE_OK) {
        printf("Download failed: %s\n", curl_easy_strerror(res));
        remove(output_path);
        return -1;
    }

    printf("Download completed: %s\n", output_path);
    return 0;
}

// 备用下载方法（使用wget）
static int download_file_with_wget(const char* url, const char* output_path)
{
    char command[1024];

    printf("Downloading with wget: %s\n", url);
    update_status(UPGRADE_STATUS_DOWNLOADING, "Downloading upgrade package");

    snprintf(command, sizeof(command),
             "wget -O %s '%s' > /tmp/neuron_upgrade_download.log 2>&1",
             output_path, url);

    int result = system(command);

    if (result == 0) {
        printf("Download completed: %s\n", output_path);
        return 0;
    } else {
        printf("Download failed with code: %d\n", result);
        return -1;
    }
}

// 下载文件
static int download_file(const char* url, const char* output_path)
{
    // 首先尝试使用CURL
    if (download_file_with_curl(url, output_path) == 0) {
        return 0;
    }

    printf("CURL download failed, trying wget...\n");

    // 如果CURL失败，尝试使用wget
    return download_file_with_wget(url, output_path);
}

// 执行升级脚本
static int execute_upgrade_script(const char* script_path)
{
    char command[256];

    printf("Executing upgrade script: %s\n", script_path);
    update_status(UPGRADE_STATUS_RUNNING, "Executing upgrade script");

    snprintf(command, sizeof(command), "%s", script_path);

    int result = system(command);
    return result;
}

// 确定升级文件路径
static const char* get_upgrade_filename(const char* download_url)
{
    const char* ext = strrchr(download_url, '.');
    if (ext) {
        if (strcmp(ext, ".zip") == 0) {
            return "/root/upgrade.zip";
        } else if (strcmp(ext, ".so") == 0) {
            // 对于.so文件，我们需要从URL中提取文件名
            const char* filename = strrchr(download_url, '/');
            if (filename) {
                static char so_path[256];
                snprintf(so_path, sizeof(so_path), "/root/%s", filename + 1);
                return so_path;
            }
        }
    }

    // 默认情况
    return "/root/upgrade.zip";
}

upgrade_result_t execute_remote_upgrade(
    const char* mqtt_broker,
    int mqtt_port,
    const char* mqtt_username,
    const char* mqtt_password,
    const char* status_topic,
    const char* download_url,
    const char* expected_md5,
    const char* upgrade_script)
{
    // 参数验证
    if (!mqtt_broker || !download_url || !upgrade_script) {
        printf("Error: Missing required parameters\n");
        return UPGRADE_ERROR_INVALID_PARAMS;
    }

    printf("Starting remote upgrade process...\n");

    // 初始化状态上报
    if (init_status_reporting(status_topic) != 0) {
        printf("Failed to initialize status reporting\n");
        return UPGRADE_ERROR_MQTT_CONNECT_FAILED;
    }

    update_status(UPGRADE_STATUS_CONNECTING, "Connecting to MQTT broker");

    // 初始化MQTT连接
    if (init_mqtt_connection(mqtt_broker, mqtt_port, mqtt_username, mqtt_password) != 0) {
        printf("Failed to initialize MQTT connection\n");
        stop_status_reporting();
        return UPGRADE_ERROR_MQTT_CONNECT_FAILED;
    }

    // 确定升级文件路径
    const char* upgrade_file = get_upgrade_filename(download_url);
    char temp_file[256];
    snprintf(temp_file, sizeof(temp_file), "/tmp/neuron_upgrade_%d.tmp", getpid());

    // 下载文件到临时位置
    if (download_file(download_url, temp_file) != 0) {
        update_status(UPGRADE_STATUS_FAILED, "Download failed");
        stop_status_reporting();
        if (g_upgrade_ctx.mqtt_client) {
            neu_mqtt_client_close(g_upgrade_ctx.mqtt_client);
            neu_mqtt_client_free(g_upgrade_ctx.mqtt_client);
        }
        return UPGRADE_ERROR_DOWNLOAD_FAILED;
    }

    update_status(UPGRADE_STATUS_DOWNLOADED, "Download completed");
    update_status(UPGRADE_STATUS_VERIFYING, "Verifying file integrity");

    // 验证MD5（如果提供了MD5值）
    if (expected_md5 && strlen(expected_md5) > 0) {
        char* actual_md5 = calculate_file_md5(temp_file);
        if (strcmp(actual_md5, expected_md5) != 0) {
            printf("MD5 verification failed. Expected: %s, Actual: %s\n",
                   expected_md5, actual_md5);

            update_status(UPGRADE_STATUS_FAILED, "MD5 verification failed");
            stop_status_reporting();

            remove(temp_file);
            if (g_upgrade_ctx.mqtt_client) {
                neu_mqtt_client_close(g_upgrade_ctx.mqtt_client);
                neu_mqtt_client_free(g_upgrade_ctx.mqtt_client);
            }
            return UPGRADE_ERROR_MD5_MISMATCH;
        }
        printf("MD5 verification passed: %s\n", actual_md5);
    }

    update_status(UPGRADE_STATUS_VERIFIED, "File verification passed");
    update_status(UPGRADE_STATUS_STARTING, "Starting upgrade process");

    // 移动文件到最终位置
    char move_cmd[512];
    snprintf(move_cmd, sizeof(move_cmd), "mv %s %s", temp_file, upgrade_file);
    if (system(move_cmd) != 0) {
        printf("Failed to move file to final location\n");
        update_status(UPGRADE_STATUS_FAILED, "Failed to prepare upgrade file");
        stop_status_reporting();
        if (g_upgrade_ctx.mqtt_client) {
            neu_mqtt_client_close(g_upgrade_ctx.mqtt_client);
            neu_mqtt_client_free(g_upgrade_ctx.mqtt_client);
        }
        return UPGRADE_ERROR_SCRIPT_FAILED;
    }

    printf("Upgrade file prepared: %s\n", upgrade_file);

    // 执行升级脚本
    int script_result = execute_upgrade_script(upgrade_script);

    // 发布最终状态
    if (script_result == 0) {
        update_status(UPGRADE_STATUS_SUCCESS, "Upgrade completed successfully");
    } else {
        update_status(UPGRADE_STATUS_FAILED, "Upgrade script execution failed");
    }

    // 等待最终状态被发送
    sleep(2);

    // 清理资源
    stop_status_reporting();

    if (g_upgrade_ctx.mqtt_client) {
        neu_mqtt_client_close(g_upgrade_ctx.mqtt_client);
        neu_mqtt_client_free(g_upgrade_ctx.mqtt_client);
    }

    if (script_result == 0) {
        printf("Remote upgrade completed successfully\n");
        return UPGRADE_SUCCESS;
    } else {
        printf("Remote upgrade failed: script returned %d\n", script_result);
        return UPGRADE_ERROR_SCRIPT_FAILED;
    }
}
